#include <iostream>
#include <thread>
#include <mutex>
#include <chrono>
#include <functional>
#include <queue>
#include <vector>
#include <condition_variable>

using task_t = std::function<void()>;

class ThreadSafeTaskQueue
{
public:
    ThreadSafeTaskQueue() : stoped_(false) {}

public:
    // 在线程安全的情况下进行任务的push
    void push(const task_t &task)
    {
        std::unique_lock<std::mutex> lock(mutex_);
        if (!stoped_)
        {
            task_queue_.push(task);
            cv_.notify_one(); // 唤醒一个消费者
            return;
        }
        else
            std::cerr << "push failed\n";
    }
    bool pop(task_t &task)
    {
        std::unique_lock<std::mutex> lock(mutex_);
        // 等待队列中有数据或者运行结束
        cv_.wait(lock, [this]()
                 { return !task_queue_.empty() || stoped_; });
        if (stoped_ || task_queue_.empty())
            return false;
        if (!stoped_)
        {
            task = task_queue_.front();
            task_queue_.pop();
            return true;
        }
        return false;
    }
    void stop()
    {
        std::unique_lock<std::mutex> lock(mutex_);
        stoped_ = true;
        cv_.notify_all();
    }

private:
    bool stoped_;
    std::mutex mutex_;
    std::condition_variable cv_;
    std::queue<task_t> task_queue_;
};
void Producer(ThreadSafeTaskQueue &task, std::size_t ProducerID)
{
    for (int i = 1; i <= 5; i++)
    {
        auto temp = [ProducerID, i]()
        {
            std::cout << "Producer: "
                      << ProducerID
                      << ": Task "
                      << i
                      << std::endl;
        };
        task.push(temp);
        std::this_thread::sleep_for(std::chrono::milliseconds(300));
    }
    std::cout << "Producer finished" << std::endl;
}

void Consumer(ThreadSafeTaskQueue &task, std::size_t ConsumerID)
{
    task_t t;
    while (task.pop(t))
    {
        std::cout << "Consumer "
                  << ConsumerID
                  << ": Execute Task -> ";
        t();
        std::this_thread::sleep_for(std::chrono::milliseconds(500));
    }
    std::cout << "Consumer " << ConsumerID << " finished" << std::endl;
}

int main()
{
    ThreadSafeTaskQueue tstq;
    std::vector<std::thread> ProducerThreads;
    std::vector<std::thread> ConsumerThreads;

    for (size_t i = 0; i < 2; i++)
    {
        ProducerThreads.emplace_back(Producer, std::ref(tstq), i);
    }
    for (size_t i = 0; i < 3; i++)
    {
        ConsumerThreads.emplace_back(Consumer, std::ref(tstq), i);
    }

    for (auto &t : ProducerThreads)
        t.join();
    std::cout << "All Producer finished." << std::endl;
    tstq.stop();

    for (auto &t : ConsumerThreads)
        t.join();
    std::cout << "All tasks processed. Program exited normally." << std::endl;
    return 0;
}
